package sasl

// Options specifies the features of the protocol on which you're running SASL.
//
// The zero value specifies a protocol that implements all the SHOULD
// features of the SASL protocol specification.
type Options struct {
	// Set to true if the protocol cannot provide the client's initial
	// response alongside the SASL START command. This will result in
	// an extra round-trip for all client-initiated mechanisms.
	ClientCannotInitiate bool
}

// SASL provides client and server state machines to authenticate
// peers.
type SASL struct {
	options    Options
	mechanisms []Mechanism
}

// Create a new SASL state machine with the given mechanisms.
//
// For clients, the order of mechanisms is most preferred to least
// preferred. For servers, the order doesn't matter since the client
// selects the mechanisms to use.
func New(mechanisms []Mechanism, opts Options) *SASL {
	// TODO: more checks on the mechanisms passed in.
	return &SASL{opts, mechanisms}
}

// List the mechanisms registered in this state machine.
func (s *SASL) Mechanisms() []string {
	mechs := make([]string, 0, len(s.mechanisms))
	for _, m := range s.mechanisms {
		mechs = append(mechs, m.Name())
	}
	return mechs
}

// Initiate a SASL exchange as a client.
//
// Returns the initial message to be sent to the server (one of the
// Msg* types) and the client state machine for this exchange.
func (s *SASL) Client(serverMechanisms []string) (interface{}, *ClientState) {
	selected := make([]Mechanism, 0, len(s.mechanisms))
	for _, m := range s.mechanisms {
		for _, sm := range serverMechanisms {
			if m.Name() == sm {
				selected = append(selected, m)
			}
		}
	}

	st := &ClientState{
		options:    s.options,
		mechanisms: selected,
		next:       0,
		current:    nil,
	}
	return st.nextMechanism(), st
}

// ClientState is the SASL state machine for the client side of a
// single SASL exchange.
type ClientState struct {
	options    Options
	mechanisms []Mechanism
	next       int
	current    MechanismClientState
}

// Advance the client state machine by one step.
//
// Given the latest message received from the server, advance the
// internal state and return the next message to send back. Both
// messages are one of the Msg* types.
//
// If MsgDone is returned, the SASL exchange is finished and can
// advance no further. It is an error to call Next() again after
// receiving MsgDone.
func (s *ClientState) Next(in interface{}) interface{} {
	switch msg := in.(type) {
	case MsgChallenge:
		resp, err := s.current.Next([]byte(msg))
		if err != nil {
			s.current.Close()
			s.current = nil
			return MsgClientCancel{}
		}
		return MsgResponse(resp)
	case MsgClientCancelAck:
		return s.nextMechanism()
	case MsgSuccess:
		return MsgDone{true, []byte(msg)}
	case MsgFailure:
		s.current.Close()
		return s.nextMechanism()
	}
	panic("Illegal state received from server")
}

func (s *ClientState) nextMechanism() interface{} {
	for s.next < len(s.mechanisms) {
		m := s.mechanisms[s.next]
		s.next++
		s.current = m.Client()
		if m.ClientInitiated() && !s.options.ClientCannotInitiate {
			init, err := s.current.Next([]byte{})
			if err != nil {
				s.current.Close()
				continue
			}
			return MsgStart{m.Name(), init}
		} else {
			return MsgStart{m.Name(), nil}
		}
	}
	return MsgDone{false, []byte("No more mechanisms to try")}
}
